#include <iostream>
#include <thread>
#include <vector>
#include <chrono>
#include <fstream>
#include <string>

// 添加OpenCV的头文件
#include <opencv2/opencv.hpp>
#include <opencv2/core.hpp>
#include <opencv2/videoio.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/face.hpp>

#define CVUI_IMPLEMENTATION // 该定义需要在一个（且只有一个）C++ 文件中使用
#include "cvui.h"

#include "TCPServer.hpp" // TCP服务器

using namespace std;
using namespace cv;
using namespace face;

#define WINDOW_NAME "Facial recognition registration system" //  显示窗口的名字
#define PORT 8888

void saveImage(int index, const Mat &resizeImg);
void updateModel(vector<Mat> &images, vector<int> &labels, int identifier);
void processImage(int index, Mat croppedImage, vector<Mat> &images, vector<int> &labels, int identifier);
void processRegister(Mat resizeImg, vector<Mat> &images, vector<int> &labels, int identifier, string &name);
string splitFaceData(string &str, int &index);

// 添加全局变量来控制弹窗显示
bool showSuccessPopup = false;
bool showPopup = false;
chrono::system_clock::time_point popupStartTime;

int main()
{
    Mat frame = Mat(480, 900, CV_8UC3),    // 主画布
        frameCam = Mat(480, 600, CV_8UC3), // 摄像头帧
        croppedImage,                      // 截取的人脸图像
        resizeImg,                         // 调整大小后人脸图像
        faceflagImge = imread("../img/face_flag.jpg"),
        registerImg = imread("../img/Register.jpg"),
        upDateImage = imread("../img/UpDate.jpg"),
        GetImage = imread("../img/GetImg.jpg"),
        successImage = imread("../img/Rsuccess.jpg");

    int Index = 0;
    // 人脸的宽度和高度
    int w;
    int h;
    // 人脸图像的起始坐标
    int x;
    int y;

    string face_msg; // 存放人脸数据
    string name;     // 人脸名称

    bool faceFlag = false; // 人脸注册和更新窗口显示的标志

    // 人脸标识符
    int currentIdentifier = -1;

    cout << "This face recognition server" << endl;

    // 初始化摄像头对象
    VideoCapture cap(0);
    if (!cap.isOpened())
    {
        cerr << "ERROR! Unable to open camera" << endl;
        return -1;
    }
    cout << "Start grabbing" << endl
         << "Press any key to terminate" << endl;

    // 准备好人脸检测器
    CascadeClassifier cascade; // 定义一个检测器

    //  加载检测器
    if (!cascade.load("E:/xuexi/text/C++/QTCode/opencv452/etc//haarcascades/haarcascade_frontalface_alt.xml"))
    {
        cerr << "ERROR: Could not load classifier cascade" << endl;
        return -1;
    }

    // 加载人脸图像和标签
    vector<Mat> images; // 人脸图像数组
    vector<int> labels; // 人脸对应用户标识 ID

    // 创建TCP服务器对象
    TCPServer TcpServer(PORT);

    TcpServer.start(); // 启动服务器

    // 初始化cvui窗口
    cvui::init(WINDOW_NAME);

    while (1)
    {
        //  使用颜色填充画布 BGR
        frame = cv::Scalar(255, 255, 255);

        // 读取采集的数据 放入 数据帧中
        cap >> frameCam;
        if (frameCam.empty())
        {
            cerr << "ERROR! Blank frameCam grabbed" << endl;
            break;
        }
        cvui::image(frame, 0, 0, frameCam);
        cvui::image(frame, 650, 0, registerImg);

        // 构造一个数组，用于存储检测到的人脸的位置（矩形）
        std::vector<Rect> faces;

        // 检测人脸，并把检测到的结果 存入faces中
        cascade.detectMultiScale(frame, faces,
                                 1.1, 3, 0, Size(100, 100), Size(300, 300));

        // 如果检测到人脸
        if (faces.size() > 0)
        {
            // 从前面检测到的人脸数据中获得人脸出现的起始坐标
            x = faces.at(0).x;
            y = faces.at(0).y;
            // 人脸的宽度和高度
            w = faces.at(0).width;
            h = faces.at(0).height;

            rectangle(frame, Point(x, y), Point(x + w, y + h),
                      Scalar(0, 255, 0), 2, LINE_8, 0);
        }

        // 显示人脸图像 更新或注册模型
        if (faceFlag)
        {
            cvui::image(frame, 690, 150, resizeImg);
            cvui::image(frame, 690, 310, upDateImage);

            // 注册
            if (cvui::button(frame, 670, 380, "Register"))
            {
                cout << "register" << endl;

                // 判断人脸数据是否正确
                if (name.length() < 1 || face_msg.length() < 3)
                {
                    cout << "name or face_msg is too short" << endl;
                    showPopup = true;
                    popupStartTime = chrono::system_clock::now();
                    faceFlag = false;
                    continue;
                }

                // 在创建线程之前验证图像是否正确复制
                if (!resizeImg.empty())
                {
                    // 创建线程处理图像和模型更新
                    thread processingThread(processRegister, ref(resizeImg), ref(images), ref(labels), currentIdentifier, ref(name));
                    processingThread.detach(); // 分离线程
                }

                if (face_msg.find('_')) // 确保格式正确
                {
                    // 将新人脸数据写入文件
                    ofstream file("../face_data/face_data.txt", std::ios::app);
                    if (file.is_open())
                    {
                        file << face_msg << endl; // 追加内容并换行
                        file.close();
                    }
                }
            }

            // 更新
            if (cvui::button(frame, 800, 380, "Update"))
            {
                cout << "Update" << endl;

                // 判断人脸数据是否正确
                if (name.length() < 1 || face_msg.length() < 3)
                {
                    cout << "name or face_msg is too short" << endl;
                    showPopup = true;
                    popupStartTime = chrono::system_clock::now();
                    faceFlag = false;
                    continue;
                }

                // 在创建线程之前验证图像是否正确复制
                if (!resizeImg.empty())
                {
                    // 创建线程处理图像和模型更新
                    thread processingThread(processImage, Index++, ref(resizeImg),
                                            ref(images), ref(labels), ref(currentIdentifier));
                    processingThread.detach(); // 分离线程
                }
            }
        }
        else
            cvui::image(frame, 660, 150, GetImage);

        // 注册成功弹窗显示
        if (showSuccessPopup)
        {
            cvui::image(frame, 640, 120, successImage);

            // 检查是否需要关闭弹窗（1秒后）
            auto currentTime = chrono::system_clock::now();
            auto elapsedTime = chrono::duration_cast<chrono::milliseconds>(
                                   currentTime - popupStartTime)
                                   .count();

            if (elapsedTime > 1000)
            { // 1000毫秒
                showSuccessPopup = false;
                faceFlag = false;
            }
        }
        if (showPopup) // 数据有误弹窗显示
        {
            cvui::image(frame, 650, 120, faceflagImge);

            // 检查是否需要关闭弹窗（3秒后）
            auto currentTime = chrono::system_clock::now();
            auto elapsedTime = chrono::duration_cast<chrono::milliseconds>(
                                   currentTime - popupStartTime)
                                   .count();

            if (elapsedTime > 3000)
            { // 3000毫秒
                showPopup = false;
                faceFlag = false;
            }
        }

        // 数据输入窗口
        cvui::window(frame, 640, 410, 260, 70, "Face data");
        cvui::input(frame, 660, 440, 180, "File name", face_msg, 0.5);
        if (cvui::button(frame, 845, 440, "Yes"))
        {
            // cout << "face msg:" << face_msg << endl;
            name = splitFaceData(face_msg, currentIdentifier);
            if (name.length() > 0)
                cout << "name:" << name << endl;
            else
            {
                showPopup = true;
                popupStartTime = chrono::system_clock::now();
            }

            // cout << "face msg:" << face_msg << endl;
        }

        // 更新组件
        cvui::update();

        // 设置窗口的名字并显示窗口
        imshow(WINDOW_NAME, frame);

        int key = waitKey(5);
        // 按下特定键（Enter键，即键码13）
        if (key == 13 && faces.size() > 0)
        {
            // 广播信息测试
            // TcpServer.broadcastMessage("data from server");
            // 截图保存
            Rect roi(x, y, w, h); // 设置截取的起始位置以及尺寸
            // 截取图像人脸部分
            frame(roi).copyTo(croppedImage);

            // 重新调整图像大小 （统一尺寸）
            resize(croppedImage, resizeImg, Size(160, 160));

            faceFlag = true;
        }
        else if (key == 27) // 按下ESC键退出
        {
            cout << "ESC key pressed. Exiting..." << endl;
            TcpServer.stop();
            break;
        }
    }

    return 0;
}

// 保存图像
void saveImage(int index, const Mat &resizeImg)
{

    // 保存截取的图像
    char FileName[32] = {0};
    snprintf(FileName, 32, "../face_data/img/%d.jpg", index);
    imwrite(FileName, resizeImg);
}

// 更新模型
void updateModel(vector<Mat> &images, vector<int> &labels, int identifier)
{
    ifstream file("../face_data/face_data.txt");
    string line;
    stringstream ss;

    // 将label转换为字符串
    string labelStr = std::to_string(identifier);
    if (file.is_open())
    {
        while (getline(file, line))
        {
            // 检查行的第一个字符是否与标签对应
            if (line.find(labelStr) == 0)
            {
                // 找到'_'的位置
                size_t pos = line.find('_');
                if (pos != string::npos)
                {
                    // 提取'_'后面的字符串
                    string userLabel = line.substr(pos + 1);
                    userLabel.erase(userLabel.find_last_not_of("\n") + 1);
                    ss << "../model/" << userLabel << ".xml";
                    break; // 找到匹配的标签后退出循环
                }
            }
        }
        file.close();
    }
    else
    {
        cout << "User Label: File not found or unable to open" << endl;
    }

    // 检查ss是否为空(文件是否有对应人脸数据)
    if (ss.str().length() < 1)
    {
        cout << "The current face is not registered" << endl;
        return;
    }

    labels.push_back(identifier); // 同一个人使用同一个标识

    // 读取得到现有文件中的训练数据
    Ptr<LBPHFaceRecognizer> recognizer = LBPHFaceRecognizer::create();
    recognizer->read(ss.str());

    // 更新现有模型（添加新用户）
    recognizer->update(images, labels);

    // 保存模型
    recognizer->save(ss.str());

    cout << "update model success" << endl;
}

// 处理图像和模型更新
void processImage(int index, Mat resizeImg, vector<Mat> &images, vector<int> &labels, int identifier)
{
    saveImage(index, resizeImg);
    Mat image = imread("../face_data/img/" + to_string(index) + ".jpg", IMREAD_GRAYSCALE);
    images.push_back(image);                 // 添加到人脸图像数组
    updateModel(images, labels, identifier); // 更新模型
}

// 处理图像和模型注册
void processRegister(Mat resizeImg, vector<Mat> &images, vector<int> &labels, int identifier, string &name)
{
    int index = 1;
    saveImage(index, resizeImg);
    Mat image = imread("../face_data/img/" + to_string(index) + ".jpg", IMREAD_GRAYSCALE);
    images.push_back(image);      // 添加到人脸图像数组
    labels.push_back(identifier); // 同一个人使用同一个标识
    // 创建 LBPHFaceRecognizer 对象
    Ptr<LBPHFaceRecognizer> recognizer = LBPHFaceRecognizer::create();
    // 训练一个全新模型
    recognizer->train(images, labels);
    // 保存模型
    recognizer->save("../model/" + name + ".xml");

    cout << "register model success" << endl;

    // 触发弹窗显示
    showSuccessPopup = true;
    popupStartTime = chrono::system_clock::now();
}

// 将名字和人脸标签分离获取
string splitFaceData(string &str, int &index)
{
    if (str.length() < 3)
        return "";

    if (str.at(0) >= '0' && str.at(0) <= '9')
    {

        index = stoi(str);
        size_t pos = str.find("#");
        if (pos != string::npos)
        {
            string name = str.substr(pos + 1);
            // cout << "index: " << index << "   name: " << name << endl;

            if (name.length() > 0) // 确保姓名不为空
            {
                str.at(pos) = '_'; // 修改人脸数据格式

                return name;
            }
        }
    }
    return "";
}
